
#include "MutualInformation.h"
#include "config.h"
#include <math.h>
#include <string.h>
#include <stdlib.h> 
#include <stdio.h>

int *GfirstStateCounts;
int *GsecondStateCounts;
int *GjointStateCounts;
float *GfirstStateProbs;
float *GsecondStateProbs;
float *GjointStateProbs;

void MemInit() { 
	GfirstStateCounts = (int *) calloc(8*FIRST,sizeof(int));
	GsecondStateCounts = (int *) calloc(8*SECND,sizeof(int));
	GjointStateCounts = (int *) calloc(8*JOINT,sizeof(int));
  
	GfirstStateProbs = (float*) calloc(8*FIRST,sizeof(float));
	GsecondStateProbs = (float*) calloc(8*SECND,sizeof(float));
	GjointStateProbs = (float*) calloc(8*JOINT,sizeof(float));
}

JointProbabilityState calculateJointProbability(short *firstVector, short *secondVector, int vectorLength, const int threadIdx)
{
  int *firstStateCounts = GfirstStateCounts + threadIdx * FIRST;
  int *secondStateCounts = GsecondStateCounts + threadIdx * SECND;
  int *jointStateCounts = GjointStateCounts+ threadIdx * JOINT;
  float *firstStateProbs = GfirstStateProbs + threadIdx * FIRST;
  float *secondStateProbs = GsecondStateProbs + threadIdx * SECND;
  float *jointStateProbs = GjointStateProbs + threadIdx * JOINT;

  int firstNumStates;
  int secondNumStates;
  int jointNumStates;
  int i;
  float length = (float) vectorLength;
  JointProbabilityState state;

  firstNumStates = normaliseArrayFast(firstVector, vectorLength);
  secondNumStates = normaliseArrayFast(secondVector, vectorLength);
  jointNumStates = firstNumStates * secondNumStates;
  
  /* Optimised for number of FP operations now O(states) instead of O(vectorLength) */
  for (i = 0; i < vectorLength; i++)
  {
    firstStateCounts[firstVector[i]] += 1;
    secondStateCounts[secondVector[i]] += 1;
    jointStateCounts[secondVector[i] * firstNumStates + firstVector[i]] += 1;
  }
  
  for (i = 0; i < firstNumStates; i++)
  {
    firstStateProbs[i] = firstStateCounts[i] / length;
  }
  
  for (i = 0; i < secondNumStates; i++)
  {
    secondStateProbs[i] = secondStateCounts[i] / length;
  }
  
  for (i = 0; i < jointNumStates; i++)
  {
    jointStateProbs[i] = jointStateCounts[i] / length;
  }

  /*free(firstStateCounts);
  free(secondStateCounts);
  free(jointStateCounts);*/

  memset(firstStateCounts, 0, FIRST*sizeof(int));
  memset(secondStateCounts, 0, SECND*sizeof(int));
  memset(jointStateCounts, 0, JOINT*sizeof(int));
  
  state.jointProbabilityVector = jointStateProbs;
  state.numJointStates = jointNumStates;
  state.firstProbabilityVector = firstStateProbs;
  state.numFirstStates = firstNumStates;
  state.secondProbabilityVector = secondStateProbs;
  state.numSecondStates = secondNumStates;

  return state;
}


float calculateMutualInformation(short *dataVector, short *targetVector, const int vectorLength, const int threadIdx)
{
  float mutualInformation = 0.0;
  int firstIndex,secondIndex;
  int i;
  JointProbabilityState state = calculateJointProbability(dataVector,targetVector,vectorLength,threadIdx);
	
 
  //printf("Taking input vectors: \n");
  //for (__int64 i = 0; i < vectorLength; ++i) printf("%d vs. %d\n", dataVector[i], targetVector[i]);

  
  //printf("NumFirstStates: %d\n", state.numFirstStates);

  /*
  ** I(X;Y) = sum sum p(xy) * log (p(xy)/p(x)p(y))
  */
  for (i = 0; i < state.numJointStates; i++)
  {
	firstIndex = i % state.numFirstStates;
	secondIndex = i / state.numFirstStates;
	
	if ((state.jointProbabilityVector[i] > 0) &&
		(state.firstProbabilityVector[firstIndex] > 0) &&
		(state.secondProbabilityVector[secondIndex] > 0))
	{
	  mutualInformation += state.jointProbabilityVector[i] *
		  log(state.jointProbabilityVector[i] / state.firstProbabilityVector[firstIndex] /
		  state.secondProbabilityVector[secondIndex]);
	}
  }
  
  //printf("Getting MI1: %d\n", mutualInformation);
  mutualInformation /= log(2.0f);
  
  //printf("Getting MI2: %d\n", mutualInformation);
  
/*  free(state.firstProbabilityVector);
  free(state.secondProbabilityVector);
  free(state.jointProbabilityVector);*/
  
  memset(state.firstProbabilityVector, 0, FIRST*sizeof(float));
  memset(state.secondProbabilityVector, 0, SECND*sizeof(float));
  memset(state.jointProbabilityVector, 0, JOINT*sizeof(float));
  
  return mutualInformation;
}

short normaliseArrayFast(short* vector, int vectorLength)
{
	short maxVal = vector[0];
	
	for (int i = 1; i < vectorLength; ++i)    
		if (vector[i] > maxVal) maxVal = vector[i];
		
	return maxVal + 1;
}
